#include "util/pool/pool.h"

namespace hawking {
namespace indexlib {

MemoryChunk Pool::dummy_chunk_;

SimpleAllocatePolicy::SimpleAllocatePolicy(
    ChunkAllocatorBase* allocator, size_t chunk_size, bool own_allocator)
    : allocator_(allocator)
    , own_allocator_(own_allocator)
    , chunk_header_(nullptr)
    , current_chunk_(nullptr)
    , chunk_size_(chunk_size)
    , use_bytes_(0)
    , total_bytes_(0) {}

SimpleAllocatePolicy::SimpleAllocatePolicy(size_t chunk_size)
    : allocator_(new SimpleAllocator())
    , own_allocator_(true)
    , chunk_header_(nullptr)
    , current_chunk_(nullptr)
    , chunk_size_(chunk_size)
    , use_bytes_(0)
    , total_bytes_(0) {}

SimpleAllocatePolicy::~SimpleAllocatePolicy() {
    Release();

    if (own_allocator_) {
        delete allocator_;
    }
    allocator_ = nullptr;
}

MemoryChunk* SimpleAllocatePolicy::Allocate(size_t n_bytes) {
    // alloc at least chunk_size_
    size_t n_alloc_size = n_bytes + sizeof(ChainedMemoryChunk);
    if (n_alloc_size <= chunk_size_) {
        n_alloc_size = chunk_size_;
    }

    ChainedMemoryChunk* next_chunk = current_chunk_ ? current_chunk_->next() : chunk_header_;
    if (next_chunk && n_alloc_size <= chunk_size_) {
        current_chunk_ = next_chunk;
    } else {
        // alloc n_alloc_size bytes by mmap
        const auto allocated_chunk = allocator_->Allocate(n_alloc_size);
        if (!allocated_chunk) {
            return nullptr;
        }
        next_chunk = new (allocated_chunk) ChainedMemoryChunk(n_alloc_size);
        if (!next_chunk) {
            return nullptr;
        }

        total_bytes_ += n_alloc_size;
        if (!chunk_header_) {
            current_chunk_ = chunk_header_ = next_chunk;
        } else {
            if (current_chunk_) {
                next_chunk->next() = current_chunk_->next();
                if (current_chunk_->next()) {
                    current_chunk_->next()->prev() = next_chunk;
                }
                current_chunk_->next() = next_chunk;
                next_chunk->prev() = current_chunk_;
            } else {
                next_chunk->next() = chunk_header_;
                chunk_header_->prev() = next_chunk;
                chunk_header_ = next_chunk;
            }
            current_chunk_ = next_chunk;
        }
    }

    use_bytes_ += current_chunk_->TotalBytes();
    return current_chunk_;
}

size_t SimpleAllocatePolicy::Release() {
    ChainedMemoryChunk* p_chunk = chunk_header_;
    ChainedMemoryChunk* p_chunk2 = nullptr;
    while (p_chunk) {
        p_chunk2 = p_chunk;
        p_chunk = p_chunk2->next();
        allocator_->Deallocate((void*)p_chunk2, p_chunk2->TotalBytes());
    }
    chunk_header_ = current_chunk_ = nullptr;
    size_t total_bytes = total_bytes_;
    use_bytes_ = total_bytes_ = 0;
    return total_bytes;
}

size_t SimpleAllocatePolicy::Reset() {
    size_t total_bytes = total_bytes_;
    if (current_chunk_ == nullptr) {
        return total_bytes;
    }
    for (ChainedMemoryChunk* chunk = chunk_header_; chunk;) {
        if (chunk->TotalBytes() <= chunk_size_) {
            chunk->Reset();
            chunk = chunk->next();
        } else {
            ChainedMemoryChunk* prevChunk = chunk->prev();
            ChainedMemoryChunk* next_chunk = chunk->next();
            total_bytes_ -= chunk->TotalBytes();
            allocator_->Deallocate((void*)chunk, chunk->TotalBytes());
            chunk = next_chunk;
            if (prevChunk) {
                prevChunk->next() = next_chunk;
            } else {
                chunk_header_ = next_chunk;
            }
            if (next_chunk) {
                next_chunk->prev() = prevChunk;
            }
        }
    }
    use_bytes_ = 0;
    current_chunk_ = nullptr;
    return total_bytes;
}

void SimpleAllocatePolicy::Clear() {
    ChainedMemoryChunk* p_chunk = chunk_header_;
    while (p_chunk) {
        p_chunk->Clear();
        p_chunk = p_chunk->next();
    }
}

bool SimpleAllocatePolicy::IsInChunk(const void* ptr) const {
    ChainedMemoryChunk* chunk = chunk_header_;
    while (chunk) {
        if (chunk->IsInChunk(ptr)) {
            return true;
        }
        chunk = chunk->next();
    }
    return false;
}

Pool::Pool(AllocatePolicy* allocate_policy, size_t align_size) :
    align_size_(align_size), mem_chunk_(&(Pool::dummy_chunk_)),
    alloc_size_(0), alloc_policy_(allocate_policy) {}

Pool::Pool(ChunkAllocatorBase* allocator, size_t chunk_size, size_t align_size) :
    Pool(new SimpleAllocatePolicy(allocator, chunk_size), align_size) {}

Pool::Pool(size_t chunk_size, size_t align_size) :
    Pool(new SimpleAllocatePolicy(chunk_size), align_size) {}

void* Pool::AllocateUnsafe(size_t n_bytes) {
    size_t alloc_size = AlignBytes(n_bytes, align_size_);
    void* ptr = mem_chunk_->Allocate(alloc_size);
    if (!ptr) {
        MemoryChunk* chunk = alloc_policy_->Allocate(alloc_size);
        if (!chunk) {
            return nullptr;
        }

        mem_chunk_ = chunk;
        ptr = mem_chunk_->Allocate(alloc_size);
    }

    alloc_size_ += alloc_size;
    return ptr;
}

void* Pool::AllocateUnsafe(size_t n_bytes, size_t alignment) {
    size_t alloc_size = AlignBytes(n_bytes, align_size_);
    void* ptr = mem_chunk_->Allocate(alloc_size, alignment);
    if (!ptr) {
        MemoryChunk* chunk = alloc_policy_->Allocate(alloc_size);
        if (!chunk) {
            return nullptr;
        }

        mem_chunk_ = chunk;
        ptr = mem_chunk_->Allocate(alloc_size, alignment);
    }

    alloc_size_ += alloc_size;
    return ptr;
}

}
}
